package utils

import (
	"fmt"
	"github.com/go-resty/resty/v2"
	jsoniter "github.com/json-iterator/go"
	"github.com/tidwall/gjson"
	"log/slog"
	"net/http"
	"time"
)

var json = jsoniter.ConfigCompatibleWithStandardLibrary

func newClient() *resty.Client {
	// 创建 Resty 客户端
	client := resty.New()
	// 设置连接池参数
	client.SetTimeout(15 * time.Second) // 设置超时时间
	client.SetHeader("Content-Type", "application/json")
	client.SetDisableWarn(true)
	client.SetTransport(&http.Transport{
		MaxIdleConns:        100,              // 设置最大空闲连接数
		MaxIdleConnsPerHost: 20,               // 设置每个主机的最大空闲连接数
		IdleConnTimeout:     30 * time.Second, // 设置空闲连接超时时间
	})
	return client
}

type EsClient struct {
	client   *resty.Client
	address  string
	num      int
	index    string
	queryMap map[string]interface{}
}

type EsDoc struct {
	Index  string                 `json:"_index"`
	Id     string                 `json:"_id"`
	Source map[string]interface{} `json:"_source"`
}

func NewEsClient(address string, username string, password string, index string, num int, query []byte) *EsClient {
	client := newClient()
	client.SetBasicAuth(username, password)
	var queryMap = make(map[string]interface{})
	if len(query) > 0 {
		err := json.Unmarshal(query, &queryMap)
		if err != nil {
			slog.Error("query 参数有误!!!")
		}
	}
	return &EsClient{
		address:  address,
		client:   client,
		num:      num,
		index:    index,
		queryMap: queryMap,
	}
}

func (s *EsClient) FirstGetSourceIdAndData() (scrollId string, data []*EsDoc) {
	body := map[string]interface{}{
		"size": s.num,
	}
	if len(s.queryMap) > 0 {
		body["query"] = s.queryMap
	}
	response, err := s.client.NewRequest().
		SetBody(body).Post(fmt.Sprintf("%s/%s/_search?scroll=1m", s.address, s.index))
	if err != nil {
		slog.Error(fmt.Sprintf("读取数据失败：%s", err))
	}
	scrollId = gjson.Get(response.String(), "_scroll_id").String()
	hits := gjson.Get(response.String(), "hits.hits").Array()
	for _, item := range hits {
		doc := new(EsDoc)
		_ = json.Unmarshal([]byte(item.String()), &doc)
		data = append(data, doc)
	}
	return
}
func (s *EsClient) GetSourceIdAndData(scrollId1 string) (scrollId string, data []*EsDoc) {
	response, err := s.client.NewRequest().SetBody(map[string]string{"scroll": "1m", "scroll_id": scrollId1}).
		Post(fmt.Sprintf("%s/_search/scroll", s.address))
	if err != nil {
		slog.Error(fmt.Sprintf("读取数据失败：%s", err))
	}
	scrollId = gjson.Get(response.String(), "_scroll_id").String()
	hits := gjson.Get(response.String(), "hits.hits").Array()
	for _, item := range hits {
		doc := new(EsDoc)
		_ = json.Unmarshal([]byte(item.String()), &doc)
		data = append(data, doc)
	}
	return
}
func (s *EsClient) GetCount() int {
	countUrl := s.address + "/" + s.index + "/" + "_count"
	body := map[string]interface{}{}
	if len(s.queryMap) > 0 {
		body["query"] = s.queryMap
	}
	res, _ := s.client.NewRequest().SetBody(body).Post(countUrl)
	return int(gjson.Get(res.String(), "count").Int())
}

func (s *EsClient) EsBulkAction(body interface{}) (*resty.Response, error) {
	return s.client.NewRequest().SetBody(body).Post(s.address + "/_bulk")
}
func (s *EsClient) Ping(pre string) bool {
	res, err := s.client.NewRequest().Get(s.address)
	if err != nil {
		return false
	}
	resStr := res.String()
	version := gjson.Get(resStr, "version.number").String()
	if version != "" {
		slog.Info(fmt.Sprintf("[%s] ElasticSearch Version : %s", pre, version))
	} else {
		slog.Error(fmt.Sprintf("[%s] 连接ES失败，请检查账号密码!!!", pre))
		return false
	}
	return true
}

// DeleteMapFields 删除 map 中指定的字段
func DeleteMapFields(fs []string, m map[string]interface{}) map[string]interface{} {
	for _, f := range fs {
		delete(m, f)
	}
	return m
}
